from Bio.Align import substitution_matrices
from .align import SubstMatrix, Table

class SubstMatrix_Bio(SubstMatrix):
    def __init__(self,sm_bio_type:str="BLOSUM62"):
        self.sm_table = substitution_matrices.load(sm_bio_type)
        self.alphabet = list(self.sm_table.alphabet)
        self.sm = dict()
        self.build()
    
    def build(self):
        for i,from_word in enumerate(self.alphabet):
            for j,to_word in enumerate(self.alphabet):
                self.sm[from_word+to_word] = int(self.sm_table[i,j])

    def get_dict(self):
        return self.sm



if __name__ == "__main__":

    seq1 = "PHSWG"
    seq2 = "HGWAG"
    gap = -8
    table = Table(seq1,seq2)
    sm = SubstMatrix_Bio()
    table.needleman_Wunsch(sm,gap)

    max_cells = table.smith_Waterman(sm,gap)
    print(table.recover_align_local())

    for seq in table.recover_align():
        print(seq)
    print("score, max {}".format(max_cells[0].value))
    for line in table.get_score():
        print("\t".join(map(str,line)))
    print("trace, start {}".format(max_cells[0].ij))
    for line in table.get_trace_back():
        print("\t".join(map(str,line)))